🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22763

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit a518667 with merge base 08ce4df ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

@claude review this code

Claude finished @metascroy's task in 4m 4s —— View job

Code Review: Split cache geometry from runtime config

I read the split (cache.h), the reader/writer (model_metadata.h/.py, constants.h), the cache implementations, the call sites, and the tests. This is a clean, well-scoped refactor — the CacheGeometry / CacheConfig boundary is coherent, the overloaded valid() set is tidy, and the test coverage (especially the six malformed-geometry fixtures round-tripped through a real .pte) is excellent. A few observations, none blocking.

Things done well

• Dropping the "layers list of size 1 broadcasts to every layer" special case removes the cfg.layers.size() == 1 ? front() : [l] branch that was duplicated in every byte layer (MLXCellCache, MLXSequenceCache, print_cache_summary). Real simplification.
• read_cache_geometry and the Python write_cache_geometry validate the same invariants independently (nonempty, equal lengths, positive heads/dims, nonnegative windows), so a malformed program is caught on both the producer and consumer side.
• op_test_runner.cpp:251 — the explicit fields[1] <= 0 guard before layers.assign(fields[1], ...) is a good catch; a negative count would otherwise wrap to a huge size_type before valid() ever ran.

backends/mlx/llm/cache.py — this quietly fixes a latent bug

shared = getattr(text_config, "num_kv_shared_layers", 0)
if shared:
    layer_types = layer_types[:-shared]

The old if hasattr(...): layer_types = layer_types[:-text_config.num_kv_shared_layers] would evaluate layer_types[:-0] → layer_types[:0] → empty list whenever num_kv_shared_layers was present and 0. The if shared: guard sidesteps that. Worth a line in the PR description since it's a behavior change beyond the mechanical refactor.

Minor: redundant valid() in module_executor.cpp:241

cache::CacheConfig cfg{max_sessions * max_session_tokens, kv_dtype};
if (initial_capacity >= 0) { cfg.initial_capacity = initial_capacity; }
if (!cache::valid(*geometry, cfg)) { ... return Error::InvalidProgram; }
...
auto built = cache::CacheFactory::global().build(backend_id, cache_kind, *geometry, cfg);

CacheFactory::build already calls valid(geometry, cfg) (cache_registry.cpp:94) and returns InvalidArgument. The check here is harmless and gives a distinct error code/message, but it is a duplicate validation — worth a comment noting the intent (clearer diagnostic) or dropping it.

Design note: policy lives inside "immutable" geometry but is mutated at runtime

cache.h:148 calls LayerGeometry "Immutable per-layer architecture facts + cache policy," yet apply_window_override in run_llm_hf.cpp takes CacheGeometry& and rewrites geometry.layers[l].policy before the cache is built. That's fine as designed — the exported window is a default a runner may override — but the "immutable" phrasing in the comment is slightly aspirational given a mutable runtime override path exists. Consider softening the comment to "model-published defaults" or similar so the override isn't surprising to a later reader.

Nits

• model_metadata.py:90 write_cache_geometry returns dict[str, object] (tensors + an int) while the sibling writers return dict[str, int]; export_model_metadata.py's all_methods was correctly widened to dict[str, object] to match. Consistent, just noting the type widened.
• read_cache_geometry requires get_windows to always be present. That's guaranteed by write_cache_geometry always emitting all three tensors, so no compatibility gap today — just flagging the coupling in case a producer ever emits geometry without windows.

Overall this looks correct and ready; the items above are polish, not correctness blockers.
• branch cache-geometry